﻿#include "AudioProcessor.h"
#include "MLProgressbar.h"

#include <iostream>

void AudioProcessor::clear() {
    avcodec_free_context(&codec_ctx);
    avformat_free_context(audio_input_ctx);
    swr_free(&swr_ctx);
}

void AudioProcessor::generate(
    std::string audio_input_path,
    AudioFormat format,
    double delay_secs,
    int sample_rate,
    int aac_bit_rate,
    int flac_bit_width,
    int flac_compression_level
) {
    std::cout << "正在编码音频……（所需时间较长，约为音频全长的20%）" << std::endl;
    _open_audio(audio_input_path);
    switch (format) {
    case AudioFormat::AAC: _set_aac_codec(sample_rate, aac_bit_rate); break;
    case AudioFormat::FLAC: _set_flac_codec(flac_bit_width, sample_rate, flac_compression_level); break;
    default: throw "未知的音频格式。";
    }
    _set_resampler();
    _convert_audio(delay_secs, sample_rate);
    _has_packet = true;
    std::cout << "音频编码完成。" << std::endl << std::endl;
}

bool AudioProcessor::has_packet() const {
    return _has_packet;
}

int AudioProcessor::sample_rate() const {
    return audio_codecpar->sample_rate;
}

std::queue<std::pair<AVPacket, int64_t>> AudioProcessor::get_packets() {
    return audio_packet_queue;
}

void AudioProcessor::_open_audio(std::string audio_input_path) {
    if (avformat_open_input(&audio_input_ctx, audio_input_path.c_str(), nullptr, nullptr) != 0) {
        throw "无法打开音频文件";
    }
    if (avformat_find_stream_info(audio_input_ctx, nullptr) < 0) {
        avformat_close_input(&audio_input_ctx);
        throw "无法获取音频流信息";
    }

    // 查找音频流
    audio_stream_index = -1;
    for (unsigned i = 0; i < audio_input_ctx->nb_streams; i++) {
        if (audio_input_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            audio_stream_index = i;
            break;
        }
    }
    if (audio_stream_index == -1) {
        avformat_close_input(&audio_input_ctx);
        throw "未找到音频流";
    }
    audio_codecpar = audio_input_ctx->streams[audio_stream_index]->codecpar;
}

void AudioProcessor::_set_aac_codec(int sample_rate, int bit_rate) {
    // 查找AAC编码器 (优先尝试libfdk_aac)
    const AVCodec* codec = avcodec_find_encoder_by_name("libfdk_aac");
    if (!codec) {
        // 回退到FFmpeg内置AAC编码器
        codec = avcodec_find_encoder(AV_CODEC_ID_AAC);
        if (!codec) {
            throw "未找到AAC编码器。";
        }
    }

    // 创建编码器上下文
    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx) {
        throw "无法分配AAC编码器上下文。";
    }

    // 配置编码参数
    codec_ctx->sample_rate = sample_rate;
    codec_ctx->bit_rate = bit_rate * 1024;  // 转换为bps
    codec_ctx->sample_fmt = AV_SAMPLE_FMT_FLTP;  // AAC常用浮点平面格式

    int nb_channels = audio_codecpar->ch_layout.nb_channels;
    if (nb_channels == 1) {
        codec_ctx->ch_layout = AV_CHANNEL_LAYOUT_MONO;
    }
    else if (nb_channels == 2) {
        codec_ctx->ch_layout = AV_CHANNEL_LAYOUT_STEREO;
    }
    else {
        throw "不受支持的声道数。请考虑使用其他音频流格式。";
    }

    // 设置编码规格为MPEG-4 AAC-LC (网络传输常用)
    if (codec->id == AV_CODEC_ID_AAC) {
        codec_ctx->profile = FF_PROFILE_AAC_LOW;
    }

    // 设置时间基 (1/采样率)
    codec_ctx->time_base = AVRational{ 1, sample_rate };

    // 打开编码器
    int ret = avcodec_open2(codec_ctx, codec, nullptr);
    if (ret < 0) {
        avcodec_free_context(&codec_ctx);
        throw "无法打开AAC编码器。";
    }
}

void AudioProcessor::_set_flac_codec(int flac_bit_width, int flac_sample_rate, int flac_compression_level) {
    const AVCodec* flac_codec = avcodec_find_encoder(AV_CODEC_ID_FLAC);
    if (!flac_codec) {
        avformat_close_input(&audio_input_ctx);
        throw "找不到FLAC编码器";
    }

    codec_ctx = avcodec_alloc_context3(flac_codec);
    if (!codec_ctx) {
        avformat_close_input(&audio_input_ctx);
        throw "无法分配FLAC编码器上下文";
    }

    codec_ctx->sample_rate = flac_sample_rate;
    codec_ctx->ch_layout = audio_codecpar->ch_layout;

    // 设置FLAC位宽
    codec_ctx->bits_per_raw_sample = flac_bit_width;
    codec_ctx->compression_level = flac_compression_level;
    switch (flac_bit_width) {
    case 16: codec_ctx->sample_fmt = AV_SAMPLE_FMT_S16; break;
    case 24: codec_ctx->sample_fmt = AV_SAMPLE_FMT_S32; break;
    case 32: codec_ctx->sample_fmt = AV_SAMPLE_FMT_S32; break;
    default:
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&audio_input_ctx);
        throw "不支持的FLAC位宽";
    }
    if (avcodec_open2(codec_ctx, flac_codec, nullptr) < 0) {
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&audio_input_ctx);
        throw "无法打开FLAC编码器";
    }
}

void AudioProcessor::_set_resampler() {
	swr_ctx = NULL;
    swr_alloc_set_opts2(
        &swr_ctx, 
        &codec_ctx->ch_layout, 
        codec_ctx->sample_fmt, 
        codec_ctx->sample_rate, 
        &audio_codecpar->ch_layout, 
        (AVSampleFormat)audio_codecpar->format, 
        audio_codecpar->sample_rate, 
        0, NULL
    );
    if (!swr_ctx) {
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&audio_input_ctx);
        throw "无法分配重采样器。";
    }

    if (swr_init(swr_ctx) < 0) {
        swr_free(&swr_ctx);
        avcodec_free_context(&codec_ctx);
        avformat_close_input(&audio_input_ctx);
        throw "无法初始化重采样器。";
    }
}

void AudioProcessor::_convert_audio(double delay_secs, int flac_sample_rate) {
    // 创建解码器上下文 (只初始化一次)
    const AVCodec* decoder = avcodec_find_decoder(audio_codecpar->codec_id);
    if (!decoder) {
        throw "找不到音频解码器";
    }

    AVCodecContext* dec_ctx = avcodec_alloc_context3(decoder);
    if (!dec_ctx) {
        throw "无法分配解码器上下文";
    }

    if (avcodec_parameters_to_context(dec_ctx, audio_codecpar) < 0) {
        avcodec_free_context(&dec_ctx);
        throw "无法复制解码器参数";
    }

    if (avcodec_open2(dec_ctx, decoder, nullptr) < 0) {
        avcodec_free_context(&dec_ctx);
        throw "无法打开解码器";
    }

    // 初始化重采样器
    if (swr_init(swr_ctx) < 0) {
        avcodec_free_context(&dec_ctx);
        throw "无法初始化重采样器";
    }

    // 分配帧和包
    AVPacket* pkt = av_packet_alloc();
    AVFrame* frame = av_frame_alloc();
    AVFrame* swr_frame = av_frame_alloc();
    AVFrame* enc_frame = nullptr;
    AVPacket* enc_pkt = av_packet_alloc();

    // 创建音频FIFO缓冲区
    AVAudioFifo* fifo = av_audio_fifo_alloc(
        codec_ctx->sample_fmt,
        codec_ctx->ch_layout.nb_channels,
        codec_ctx->frame_size
    );

    // 时间戳管理 (只初始化一次)
    int64_t current_pts = 0;
    const int64_t delay_samples = static_cast<int64_t>(delay_secs * codec_ctx->sample_rate);

    // 添加延迟所需的静音
    if (delay_samples > 0) {
        AVFrame* silence_frame = av_frame_alloc();
        silence_frame->nb_samples = delay_samples;
        silence_frame->ch_layout = codec_ctx->ch_layout;
        silence_frame->format = codec_ctx->sample_fmt;
        silence_frame->sample_rate = codec_ctx->sample_rate;

        if (av_frame_get_buffer(silence_frame, 0) < 0) {
            av_frame_free(&silence_frame);
            throw "无法分配静音帧缓冲区";
        }

        av_samples_set_silence(
            silence_frame->data,
            0,
            delay_samples,
            codec_ctx->ch_layout.nb_channels,
            static_cast<AVSampleFormat>(codec_ctx->sample_fmt)
        );

        av_audio_fifo_write(fifo, (void**)silence_frame->data, delay_samples);
        av_frame_free(&silence_frame);
    }

    // 主处理循环
    while (av_read_frame(audio_input_ctx, pkt) >= 0) {
        if (pkt->stream_index != audio_stream_index) {
            av_packet_unref(pkt);
            continue;
        }

        // 发送包到解码器
        int ret = avcodec_send_packet(dec_ctx, pkt);
        av_packet_unref(pkt);
        if (ret < 0 && ret != AVERROR(EAGAIN)) continue;

        while (true) {
            ret = avcodec_receive_frame(dec_ctx, frame);
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) break;
            if (ret < 0) throw "音频解码错误";

            // 关键修改：不再在_encode_frame中修改current_pts
            _encode_frame(dec_ctx, frame, swr_frame, enc_frame, current_pts, fifo, enc_pkt, ret);
        }
    }

    // 刷新解码器
    avcodec_send_packet(dec_ctx, nullptr);
    while (true) {
        int ret = avcodec_receive_frame(dec_ctx, frame);
        if (ret == AVERROR_EOF) break;
        if (ret < 0 && ret != AVERROR(EAGAIN)) break;
        if (ret == AVERROR(EAGAIN)) continue;

        // 关键修改：不再在_encode_frame中修改current_pts
        _encode_frame(dec_ctx, frame, swr_frame, enc_frame, current_pts, fifo, enc_pkt, ret);
    }

    // 处理FIFO中剩余样本
    const int remaining = av_audio_fifo_size(fifo);
    if (remaining > 0) {
        enc_frame = av_frame_alloc();
        enc_frame->nb_samples = remaining;
        enc_frame->ch_layout = codec_ctx->ch_layout;
        enc_frame->format = codec_ctx->sample_fmt;
        enc_frame->sample_rate = codec_ctx->sample_rate;

        if (av_frame_get_buffer(enc_frame, 0) < 0) {
            throw "无法分配编码帧缓冲区";
        }

        if (av_audio_fifo_read(fifo, (void**)enc_frame->data, remaining) < remaining) {
            av_frame_free(&enc_frame);
            throw "FIFO读取错误";
        }

        enc_frame->pts = current_pts;
        current_pts += remaining;

        if (avcodec_send_frame(codec_ctx, enc_frame) < 0) {
            av_frame_free(&enc_frame);
            throw "无法发送帧到编码器";
        }

        av_frame_free(&enc_frame);
    }

    // 刷新编码器
    avcodec_send_frame(codec_ctx, nullptr);
    while (true) {
        av_packet_unref(enc_pkt);
        int ret = avcodec_receive_packet(codec_ctx, enc_pkt);
        if (ret == AVERROR_EOF) break;
        if (ret < 0 && ret != AVERROR(EAGAIN)) break;
        if (ret == AVERROR(EAGAIN)) continue;

        // 关键修改：不再应用延迟（延迟已通过静音帧实现）
        enc_pkt->stream_index = 1;
        enc_pkt->duration = codec_ctx->frame_size;

        // 正确设置时间戳
        enc_pkt->pts = enc_pkt->dts = current_pts;
        current_pts += enc_pkt->duration;

        // 深拷贝音频包
        AVPacket* packet_copy = av_packet_alloc();
        av_packet_ref(packet_copy, enc_pkt);
        audio_packet_queue.push({ *packet_copy, packet_copy->pts });
    }

    // 清理资源
    av_audio_fifo_free(fifo);
    av_packet_free(&pkt);
    av_frame_free(&frame);
    av_frame_free(&swr_frame);
    av_packet_free(&enc_pkt);
    avcodec_free_context(&dec_ctx);
}

void AudioProcessor::set_stream(AVStream* audio_stream, AVFormatContext* fmt_ctx) {
    avcodec_parameters_from_context(audio_stream->codecpar, codec_ctx);
    audio_stream->time_base = AVRational{ 1, codec_ctx->sample_rate };
    audio_stream->index = 1;
}

void AudioProcessor::_encode_frame(AVCodecContext* dec_ctx, AVFrame* frame, AVFrame* swr_frame, AVFrame* enc_frame, int64_t& current_pts, AVAudioFifo* fifo, AVPacket* enc_pkt, int& ret) {
    // 接收解码后的帧
    int out_samples = av_rescale_rnd(
        swr_get_delay(swr_ctx, frame->sample_rate) + frame->nb_samples,
        codec_ctx->sample_rate,
        frame->sample_rate,
        AV_ROUND_UP
    );

    // 配置重采样帧
    swr_frame->nb_samples = out_samples;
    swr_frame->ch_layout = codec_ctx->ch_layout;
    swr_frame->format = codec_ctx->sample_fmt;
    swr_frame->sample_rate = codec_ctx->sample_rate;

    if (av_frame_get_buffer(swr_frame, 0) < 0) {
        throw "无法分配重采样缓冲区";
    }

    // 执行重采样
    int samples_converted = swr_convert(
        swr_ctx,
        swr_frame->data, out_samples,
        (const uint8_t**)frame->data, frame->nb_samples
    );

    if (samples_converted < 0) throw "音频重采样失败";
    swr_frame->nb_samples = samples_converted;

    // 写入FIFO缓冲区
    av_audio_fifo_write(fifo, (void**)swr_frame->data, samples_converted);
    av_frame_unref(swr_frame);
    av_frame_unref(frame);

    // 从FIFO取出完整帧进行编码
    while (av_audio_fifo_size(fifo) >= codec_ctx->frame_size) {
        enc_frame = av_frame_alloc();
        enc_frame->nb_samples = codec_ctx->frame_size;
        enc_frame->ch_layout = codec_ctx->ch_layout;
        enc_frame->format = codec_ctx->sample_fmt;
        enc_frame->sample_rate = codec_ctx->sample_rate;

        if (av_frame_get_buffer(enc_frame, 0) < 0) {
            av_frame_free(&enc_frame);
            throw "无法分配编码帧缓冲区";
        }

        if (av_audio_fifo_read(fifo, (void**)enc_frame->data, codec_ctx->frame_size) < codec_ctx->frame_size) {
            av_frame_free(&enc_frame);
            throw "FIFO读取错误";
        }

        // 设置时间戳
        enc_frame->pts = current_pts;
        current_pts += codec_ctx->frame_size;

        // 发送到编码器
        if (avcodec_send_frame(codec_ctx, enc_frame) < 0) {
            av_frame_free(&enc_frame);
            throw "无法发送帧到编码器";
        }

        // 接收编码包
        while (true) {
            av_packet_unref(enc_pkt);
            ret = avcodec_receive_packet(codec_ctx, enc_pkt);
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) break;
            if (ret < 0) throw "接收音频包失败";

            // 关键修复：使用统一的时间基
            enc_pkt->stream_index = 1;
            enc_pkt->duration = codec_ctx->frame_size;

            // 设置 PTS 和 DTS 相同
            enc_pkt->pts = enc_pkt->dts = enc_frame->pts;

            // 深拷贝音频包
            AVPacket* packet_copy = av_packet_alloc();
            av_packet_ref(packet_copy, enc_pkt);

            // 存储原始时间戳（使用音频编码器的时间基）
            int64_t original_pts = enc_pkt->pts;
            audio_packet_queue.push({ *packet_copy, original_pts });
        }

        av_frame_free(&enc_frame);
    }
}